use std::env;
use std::io;
use std::fs;
use std::path::Path;

fn main() {
    let args: Vec<String> = env::args().collect();
    if args.len() < 2 {
        println!("usage: tree ./ [2]");
        return;
    }

    let p = args.get(1).unwrap();
    if !fs::metadata(p).unwrap().is_dir() {
        println!("{} is not a dir", p);
        return;
    }
    let mut m = 2_isize;
    if args.len() == 3 {
        m = args.get(2).unwrap().parse::<isize>().unwrap();
    }
    println!("{}", p);
    let d = Path::new(p); 
    let l = fs::read_dir(&d).unwrap().count();
    let _ = visit_dir(m, vec![], &d, l);
}

fn visit_dir(max:isize, mut df:Vec<bool>, dir: &Path, l: usize) -> io::Result<Vec<bool>> {
    let list = fs::read_dir(dir)?;
    for (i, entry) in list.enumerate() {
        let entry = entry?;
        let path = entry.path();
        let sf = i == (l - 1);
        for f in &df {
            if *f {
                print!("{}", "    ");
            }else{
                print!("{}", "│   ");
            }
        }
        if sf {
            print!("└── ");
        }else{
            print!("├── ");
        }
        print!("{}", path.file_name().unwrap().to_str().unwrap());        
        if path.is_dir() {
            let sl = fs::read_dir(&path)?.count(); 
            if (df.len() as isize) >= (max - 1) && max > 0 {
                println!(" (..{})", sl);
            }else{
                println!("");
            
                df.push(sf);
                df = visit_dir(max, df, &path, sl)?;
                df.pop();
            }
        }else{
            println!("");
        }
    }
    Ok(df)
}
